import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import tensorflow as tf 

class NN(object):
    def __init__(self):
        self.df = pd.read_csv('./timing/scenic_data.csv')
        self.df = self.df.select_dtypes(include=['number'])  # 仅保留数值列
        self.df = self.df.fillna(self.df.mean())  # 填充缺失值

    def create_dataset(self, df, n_steps):
        X, y = [], []
        for i in range(len(df) - n_steps):
            x_values = df.iloc[i:i+n_steps].copy()
            if 'count' in x_values.columns:
                x_values['count'] = 0  # 屏蔽目标列，避免数据泄露
            X.append(x_values.values)
            y.append(df.iloc[i + n_steps]['count'])  # 预测下一个时间步的count
        return np.array(X), np.array(y)

    def get_model(self):
        n_steps = 7  # 7天序列长度
        X, y = self.create_dataset(self.df, n_steps)
        
        train_size = int(len(X) * 0.8)
        X_train, X_test = X[:train_size], X[train_size:]
        y_train, y_test = y[:train_size], y[train_size:]
        
        scaler_X = StandardScaler()
        n_features = X_train.shape[2]
        X_train_reshaped = X_train.reshape(-1, n_features)
        X_test_reshaped = X_test.reshape(-1, n_features)
        X_train_scaled = scaler_X.fit_transform(X_train_reshaped).reshape(X_train.shape)
        X_test_scaled = scaler_X.transform(X_test_reshaped).reshape(X_test.shape)
        
        scaler_y = StandardScaler()
        y_train_scaled = scaler_y.fit_transform(y_train.reshape(-1, 1)).flatten()
        y_test_scaled = scaler_y.transform(y_test.reshape(-1, 1)).flatten()
        
        model = tf.keras.models.Sequential([
            tf.keras.layers.Input(shape=(n_steps, n_features)),  
            tf.keras.layers.LSTM(50, activation='relu', return_sequences=True),
            tf.keras.layers.LSTM(50, activation='relu'),
            tf.keras.layers.Dense(1)
        ])
        model.compile(optimizer='adam', loss='mse')
        
        model.fit(
            X_train_scaled, y_train_scaled,
            epochs=50,
            validation_data=(X_test_scaled, y_test_scaled)
        )
        
        loss = model.evaluate(X_test_scaled, y_test_scaled)
        print(f"测试集损失: {loss}")
        
        tf.keras.models.save_model(model, 'timing/my_model.keras')

if __name__ == '__main__':
    nm = NN()
    nm.get_model()
